"""
pip install openpyxl

python fujian1_write_result.py
"""

import json
import os
import pandas as pd
import re

# 定义输入和输出路径（相对路径）
# input_dir = "./fujian/fujian1/smooth_then_arima"
# input_dir = "./fujian/fujian1/pmdarima"
input_dir = "./fujian/fujian1/pmdarima_single_batch"


output_file = "./fujian/result.xlsx"
sheet_name = "月库存预测结果"

# 读取 Excel 文件
xls = pd.ExcelFile(output_file)

# 检查目标工作表是否存在
if sheet_name not in xls.sheet_names:
    raise ValueError(f"Sheet '{sheet_name}' not found in the Excel file.")

# 读取现有数据
existing_data = pd.read_excel(xls, sheet_name=sheet_name, index_col=0)

# 处理每一个 JSON 文件
for filename in os.listdir(input_dir):
    if filename.endswith(".json"):
        input_path = os.path.join(input_dir, filename)

        # 读取 JSON 文件
        with open(input_path, 'r') as f:
            json_data = json.load(f)

        # 打印文件名进行调试
        print(f"Processing file: {filename}")
        
        # 使用正则表达式提取 category 数字
        match = re.search(r'category(\d+)', filename)
        
        if match:
            category_index = int(match.group(1)) - 1  # 计算索引
        else:
            print(f"Warning: No valid category found in '{filename}'. Skipping.")
            continue

        # 提取 2023 年 7、8、9 月的库存数据
        for entry in json_data:
            date = entry['date']
            inventory = entry['inventory']
            if date.startswith("2023-07"):
                existing_data.at[f'category{category_index + 1}', '7月库存量'] = inventory
            elif date.startswith("2023-08"):
                existing_data.at[f'category{category_index + 1}', '8月库存量'] = inventory
            elif date.startswith("2023-09"):
                existing_data.at[f'category{category_index + 1}', '9月库存量'] = inventory

# 将更新后的数据写入指定的工作表
with pd.ExcelWriter(output_file, engine='openpyxl', mode='a', if_sheet_exists='replace') as writer:
    existing_data.to_excel(writer, sheet_name=sheet_name, index=True)

print("数据已成功更新到 Excel 工作表:", sheet_name)
